# -*- coding: utf-8 -*-

"""最小生成树算法之 Kruskal，具体见书"""

class Kruskal():
    def __init__(self) -> None:
        pass

    def mini_spanning_tree(self, vertex, edge, w):
        """input:
        vertex:= {a,b,c,d,e}
        edge:= {{a,b}, {d,c}, {b,e}, {c,e}}
        w:= f: edge -> number
        -------------------------
        vt 生成树族的顶点集族
        et 生成树族的边"""

        def findw(u):
            for w in vt:
                if u in w:
                    return w

        vt, et = set([frozenset([v]) for v in vertex]), set()
        while len(vt) > 1:
            w1, w2, e12, w12 = None, None, None, float('inf')
            for e in edge - et:    # 生成树族边的补集

                e = set(e)
                u, v = e.pop(), e.pop()
                wu, wv = findw(u), findw(v)    # 找到这条边所在的顶点集
                e.add(u)
                e.add(v)
                e = frozenset(e)

                if wu != wv and w(e) < w12:    # 如果两个顶点属于不同的顶点集 且 这条边的权要小些
                    w1, w2, e12, w12 = wu, wv, e, w(e)

            # 找到生成树与生成树间最小的边，更新顶点族集和边集
            vt.discard(w1)
            vt.discard(w2)
            vt.add(frozenset(w1 | w2))
            et.add(e12)

        return et


if __name__ == '__main__':
    vertex = {'a','b','c','d','e','f'}
    edge = {frozenset(('a','b'))
            , frozenset(('d','e'))
            , frozenset(('e','f'))
            , frozenset(('a','d'))
            , frozenset(('b','e'))
            , frozenset(('c','f'))
            , frozenset(('a','e'))
            , frozenset(('b','d'))
            , frozenset(('c','e'))}
    def w(e):
        weight = {
            frozenset(('a','b')): 1
            , frozenset(('d','e')): 4
            , frozenset(('e','f')): 7
            , frozenset(('a','d')): 4
            , frozenset(('b','e')): 2
            , frozenset(('c','f')): 5
            , frozenset(('a','e')): 3
            , frozenset(('b','d')): 3
            , frozenset(('c','e')): 4
        }
        return weight[e]
    
    kruskal = Kruskal()
    mini_tree = kruskal.mini_spanning_tree(vertex, edge, w)
    print(mini_tree)
